// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/dwio/common/seekable_input_stream.h>

namespace kumo::pollux {
    // Corresponds to declaration in LazyVector.h but repeated here to avoid
    // dependency.
    using RowSet = melon::Range<const int32_t *>;
} // namespace kumo::pollux

namespace kumo::pollux::dwio::common {
    inline void skipBytes(
        uint64_t numBytes,
        SeekableInputStream *input,
        const char *&bufferStart,
        const char *&bufferEnd) {
        // bufferStart and bufferEnd may be null if we haven't started reading yet.
        if (bufferEnd - bufferStart >= numBytes) {
            bufferStart += numBytes;
            return;
        }
        numBytes -= bufferEnd - bufferStart;
        input->SkipInt64(numBytes);
        bufferStart = bufferEnd;
    }

    inline void readBytes(
        int64_t numBytes,
        SeekableInputStream *input,
        void *bytes,
        const char *&bufferStart,
        const char *&bufferEnd) {
        char *bytesAsChars = reinterpret_cast<char *>(bytes);
        uint64_t bytesRead = 0;
        auto bytesNeeded = numBytes;
        for (;;) {
            auto copyBytes = std::min<uint64_t>(bytesNeeded, bufferEnd - bufferStart);
            if (copyBytes) {
                memcpy(&bytesAsChars[bytesRead], bufferStart, copyBytes);
                bytesNeeded -= copyBytes;
                bytesRead += copyBytes;
                bufferStart += copyBytes;
            }
            if (!bytesNeeded) {
                break;
            }
            int32_t size;
            const void *bufferPointer;
            if (!input->Next(&bufferPointer, &size)) {
                POLLUX_CHECK(false, "Reading past end");
            }
            bufferStart = static_cast<const char *>(bufferPointer);
            bufferEnd = bufferStart + size;
        }
    }

    template<typename T>
    inline bool isDense(const T *values, int32_t size) {
        return (values[size - 1] - values[0] == size - 1);
    }

    template<typename Dense, typename Sparse, typename SparseN>
    void rowLoop(
        const int32_t *rows,
        int32_t begin,
        int32_t end,
        int32_t step,
        Dense dense,
        Sparse sparse,
        SparseN sparseN) {
        int32_t i = begin;
        auto firstPartial = (end - begin) & ~(step - 1);
        for (; i < firstPartial; i += step) {
            if (isDense(&rows[i], step)) {
                dense(i);
            } else {
                sparse(i);
            }
        }
        if (i != end) {
            sparseN(i, end - i);
        }
    }

    template<typename T, typename TResult>
    inline void readContiguous(
        int32_t size,
        SeekableInputStream &input,
        TResult *output,
        const char *&bufferStart,
        const char *&bufferEnd) {
        if constexpr (sizeof(T) == sizeof(TResult)) {
            readBytes(
                size * sizeof(T),
                &input,
                reinterpret_cast<char *>(output),
                bufferStart,
                bufferEnd);
            return;
        }
        auto toRead = size;
        int32_t numOut = 0;
        while (toRead) {
            auto available = (bufferEnd - bufferStart) / sizeof(T);
            if (available == 0) {
                T temp;
                readBytes(
                    sizeof(T),
                    &input,
                    reinterpret_cast<char *>(&temp),
                    bufferStart,
                    bufferEnd);
                output[numOut++] = temp;
                --toRead;
                continue;
            }
            auto pos = bufferStart;
            auto numRead = std::min<int32_t>(available, toRead);
            for (auto i = 0; i < numRead; ++i) {
                output[numOut++] = *reinterpret_cast<const T *>(pos);
                pos += sizeof(T);
            }
            bufferStart = pos;
            toRead -= numRead;
        }
    }

    // Returns the number of elements in rows that are < limit.
    inline int32_t numBelow(melon::Range<const int32_t *> rows, int32_t limit) {
        return std::lower_bound(rows.begin(), rows.end(), limit) - rows.begin();
    }

    template<typename T, typename SingleValue, typename SparseRange>
    inline void loopOverBuffers(
        melon::Range<const int32_t *> rows,
        int32_t initialRow,
        SeekableInputStream &input,
        const char *&bufferStart,
        const char *&bufferEnd,
        SingleValue singleValue,
        SparseRange range) {
        int32_t rowOffset = initialRow;
        int32_t rowIndex = 0;
        while (rowIndex < rows.size()) {
            auto available = (bufferEnd - bufferStart) / sizeof(T);
            auto numRowsInBuffer = rows.back() - rowOffset < available
                                       ? rows.size() - rowIndex
                                       : numBelow(
                                           melon::Range<const int32_t *>(
                                               &rows[rowIndex], rows.size() - rowIndex),
                                           rowOffset + available);

            if (!numRowsInBuffer) {
                skipBytes(
                    (rows[rowIndex] - rowOffset) * sizeof(T),
                    &input,
                    bufferStart,
                    bufferEnd);
                T temp;
                readBytes(
                    sizeof(T),
                    &input,
                    reinterpret_cast<char *>(&temp),
                    bufferStart,
                    bufferEnd);
                singleValue(temp, rowIndex);
                rowOffset = rows[rowIndex] + 1;
                ++rowIndex;
                continue;
            }
            range(
                rows.data(),
                rowIndex,
                numRowsInBuffer,
                rowOffset,
                reinterpret_cast<const T *>(bufferStart));
            bufferStart +=
                    (rows[rowIndex + numRowsInBuffer - 1] + 1 - rowOffset) * sizeof(T);
            rowOffset = rows[rowIndex + numRowsInBuffer - 1] + 1;
            rowIndex += numRowsInBuffer;
        }
    }

    template<typename T, typename TResult>
    inline void readRows(
        melon::Range<const int32_t *> rows,
        int32_t initialRow,
        SeekableInputStream &input,
        TResult *output,
        const char *&bufferStart,
        const char *&bufferEnd) {
        int32_t numOut = 0;
        loopOverBuffers<T>(
            rows,
            initialRow,
            input,
            bufferStart,
            bufferEnd,
            [&](T value, int32_t /*row*/) { output[numOut++] = value; },
            [&](const int32_t *rows,
                int32_t rowIndex,
                int32_t numRowsInBuffer,
                int32_t rowOffset,
                const T *buffer) {
                for (auto i = 0; i < numRowsInBuffer; ++i) {
                    output[numOut++] = buffer[rows[i + rowIndex] - rowOffset];
                }
            });
    }
} // namespace kumo::pollux::dwio::common
